#include "log_sum_exp_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"
#include <algorithm>
#include <cmath>
#include <vector>
#include <limits>
namespace optiling {
    const uint32_t BLOCK_SIZE = 64;
static ge::graphStatus TilingFunc(gert::TilingContext* context)
{
  LogSumExpTilingData tiling;
  uint32_t sizeofdatatype;
  uint32_t totalLengthAligned;
  auto ascendcPlatform = 
      platform_ascendc::PlatformAscendC(context->GetPlatformInfo());
  auto socVersion = ascendcPlatform.GetSocVersion();
  uint64_t ub_size;
  ascendcPlatform.GetCoreMemSize(platform_ascendc::CoreMemType::UB, ub_size);
  //uto axes = context->GetAttrs()->GetListInt(0);
  auto dim = context->GetAttrs()->GetListInt(0);
  uint32_t totalLength = 1;
  uint32_t stepSize = 1;
  uint32_t unitCount = 1;
  auto axesDimNum = dim->GetSize();
  
  std::cout<<"axesDimNum:"<<axesDimNum<<std::endl;
  int32_t a =*(dim->GetData());
  std::cout<<"*(dim->GetData()):"<<a<<std::endl;
  // std::cout<<"ax"
  if(axesDimNum ==1 && (*(dim->GetData()) ==1 || *(dim->GetData()) ==(-1)) ){
    context->SetTilingKey(2);
    std::cout<<"走 dim =1的分支"<<std::endl;
    auto shape_x = context->GetInputTensor(0)->GetOriginShape();
    int64_t M = 1, N = 1, Z = 1;
    if(shape_x.GetDimNum() > 2)
    {
        for (int i = 0; i < shape_x.GetDimNum() - 2; i++)
            Z *= shape_x.GetDim(i);
        M = shape_x.GetDim(shape_x.GetDimNum() - 2);
        N = shape_x.GetDim(shape_x.GetDimNum() - 1);
        // Z = shape_x.GetDim(0);
    }
    else if(shape_x.GetDimNum() == 2)
    {
        M = shape_x.GetDim(0);
        N = shape_x.GetDim(1);
        Z = 1;
    }
    else if(shape_x.GetDimNum() == 1){
        N=shape_x.GetDim(0);
        M=1;
        Z=1;
    }
    tiling.set_M(M);
    tiling.set_N(N);
    tiling.set_Z(Z);


  }
  //(axesDimNum != 0 && *(dim->GetData()) !=1 )
  else  
  {
    
    std::cout<<"走 dim =0的分支"<<std::endl;
    context->SetTilingKey(1);
    // 处理多轴情况，axes为空时，不需要计算
    // 仅处理了axes列表中为连续维度的情况，例如[0], [1, 2]之类
    // 因为是连续维度，所以定位通过axes列表的值，将shape中的维度分为三部分
    // 前部分（unitCount），中间部分(stepSize)，后部分(totalLength)
    auto shape = context->GetInputShape(0)->GetOriginShape();
    uint32_t dimNum = shape.GetDimNum();
    uint32_t maxdim = 0;
    uint32_t mindim = dimNum;
    for (int i = 0; i < axesDimNum; i++)
    {
      auto index = *(dim->GetData() + i);
      if (index == -2){
        index =0;
      }
      if (index ==-1){
        index =1;
      }
      if (index < mindim)
        mindim = index;
      if (index > maxdim)
        maxdim = index;
    }
    if (maxdim > dimNum)
    {
      totalLength = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
      stepSize = 1;
      unitCount = 1;
    }
    else
    {
      for (int i = 0; i < dimNum; i++)
      {
        if (i < mindim)
        {
          unitCount *= shape.GetDim(i);
        }
        else if (i > maxdim)
        {
          stepSize *= shape.GetDim(i);
        }
        else
        {
          totalLength *= shape.GetDim(i);
        }
      }
    }
  }
  
  // else
  // {
  //   stepSize = 1;
  //   unitCount = 1;
  //   totalLength = context->GetInputShape(0)->GetStorageShape().GetShapeSize();
  // }
  tiling.set_totalLength(totalLength);
  tiling.set_stepSize(stepSize);
  tiling.set_unitCount(unitCount);
  auto dt = context->GetInputDesc(0)->GetDataType();
  uint32_t typeKey = 0;
  uint32_t dataNum = 0;
  // 不同类型下一次tile处理中，ub中最多存储的tileLength的倍数
  // 即一次tile处理中长度为tileLength的数据块个数
  if (dt == 1)
  {
    sizeofdatatype = 2;
    dataNum = 6;  //6
    tiling.set_typeKey(1);
  }
  else if (dt == 0)
  {
    sizeofdatatype = 4;
    dataNum = 6;  //1 2 3 4 5 6(有点戏) 7 8 9 10 11 12 13 14 15(1)  16(有点戏) 
    //17   18   19  20  21 22 23 24
    tiling.set_typeKey(0);
  }
  uint32_t ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
  uint32_t ub_block_num =
      ((ub_size) / BLOCK_SIZE / dataNum);
    // uint32_t ub_block_num =
    //   ((ub_size) / BLOCK_SIZE / dataNum);
  uint32_t tile_num;
  if (ub_block_num % 2 != 0)
  {
    ub_block_num = ub_block_num - 1;
  }
  if (totalLength % ALIGN_NUM != 0)
  {
    totalLengthAligned =
        ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
  }
  else
  {
    totalLengthAligned = totalLength;
  }
  //context->SetBlockDim(1);
  uint32_t tileLength = 0;
  uint32_t lasttileLength = 0;
  tile_num = totalLengthAligned / ALIGN_NUM / ub_block_num;
  if ((totalLengthAligned / ALIGN_NUM) % ub_block_num == 0 ||
      tile_num == 0)
  {
    if (tile_num == 0)
    {
      tile_num = 1;
    }
    if (totalLengthAligned < ub_block_num * ALIGN_NUM)
    {
      tileLength = totalLengthAligned;
      lasttileLength = tileLength;
    }
    else
    {
      tileLength = ub_block_num * ALIGN_NUM;
      lasttileLength = tileLength;
    }
  }
  else
  {
    tile_num = tile_num + 1;
    tileLength = ub_block_num * ALIGN_NUM;
    lasttileLength = totalLengthAligned - (tile_num - 1) * tileLength;
  }
  if (totalLength % tileLength)
  {
    lasttileLength = totalLength % tileLength;
  }
  tiling.set_blockLength(totalLengthAligned);
  tiling.set_tileNum(tile_num);
  tiling.set_tileLength(tileLength);
  tiling.set_lasttileLength(lasttileLength);

  

  context->SetBlockDim(1);
  tiling.SaveToBuffer(context->GetRawTilingData()->GetData(), context->GetRawTilingData()->GetCapacity());
  context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
  size_t *currentWorkspace = context->GetWorkspaceSizes(1);
  currentWorkspace[0] = 0;
  return ge::GRAPH_SUCCESS;
}
}

namespace ge {
static ge::graphStatus InferShape(gert::InferShapeContext* context)
{
    const gert::Shape* x1_shape = context->GetInputShape(0);
    gert::Shape* y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
}
}

namespace ops {
class LogSumExp : public OpDef {
public:
    explicit LogSumExp(const char* name) : OpDef(name)
    {
        this->Input("x")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Output("y")
            .ParamType(REQUIRED)
            .DataType({ge::DT_FLOAT, ge::DT_FLOAT16})
            .Format({ge::FORMAT_ND, ge::FORMAT_ND})
            .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
        this->Attr("dim").AttrType(OPTIONAL).ListInt({0});
        this->Attr("keep_dim").AttrType(OPTIONAL).Bool(false);
        this->SetInferShape(ge::InferShape);
        this->AICore()
            .SetTiling(optiling::TilingFunc);
        this->AICore().AddConfig("ascend310b");
    }
};
OP_ADD(LogSumExp);
}
